################################################################################
#       Copyright (C) 2010 Michael Yurko <myurko@gmail.com>
#
#This file is part of pyQAP.
#
#pyQAP is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, either version 2 of the License, or
#(at your option) any later version.
#
#pyQAP is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with pyQAP.  If not, see <http://www.gnu.org/licenses/>.
################################################################################

'''
Contains the classes for QAP
'''

cimport numpy as np
cimport cython

import cython
import numpy as np
import cost
import vars
from cPickle import dumps, loads

#type declarations
ctypedef Py_ssize_t index_t
ctypedef Py_ssize_t np_int_t
ctypedef Py_ssize_t int_t
ctypedef np.float64_t float_t
ctypedef Py_ssize_t bool_t

cdef class PSolution:
    """
    This class holds the partial solutions used by the solve function.
    It contains the attributes n , sol (the partial solution array), unplaced
    (a list of unplaced facilities), and placed (the number of facilities 
    currently placed. In unplaced, if unplaced[i] == True, then facility i is
    not assigned to any location in the current PSolution. If unplaced[i] ==
    False, then facility i is placed in the current PSolution. The array sol is
    the permutation vector of the PSolution. In it the assignment sol[i] = j 
    means that facility j is assigned to location i. It is important to note
    that the value of sol[i] is only meaningful for i in [0,placed) since
    the array is instantiated with whatever happens to be in the memory before
    the array is allocated. In essence, it is "random" data after sol[placed-1].
    """
    
    cdef index_t n, placed
    cdef np.ndarray sol, unplaced
    cdef np_int_t *sol_p, *unplaced_p
    
    def __init__(self, n=None, problem = None, sol = None,
                 unplaced = None, placed = None):
        """
        Initialize PSolution for the first problem instance of size n.
        
        When n is none, then an empty PSolution is created. If n is specified,
        then an initial solution of size n is created. I.e. an empty solution
        with all unplaced.copy
        
        :param n: the size of the solution
        :type n: int
        :param problem: the problem that this partial solution corresponds to
        :type problem: Problem
        :param sol: the solution array
        :type sol: ndarray of size n of integers
        :param unplaced: an array that marks which facilities are assigned
        :type unplaced: ndarray of size n of integers
        :param placed: the number of locations that have been placed
        :type placed: int
        :returns: a partial solution 
        :rtype: PSolution
        
        Examples:
        
        >>> from pyQAP import *
        >>> from classes import PSolution
        >>> print classes.PSolution(5)
        Partial QAP solution: [...] with 0 placed.
         Placed:0
         Unplaced:[1 1 1 1 1]
        >>> PSolution(10)
        Partial QAP solution: [...] with 0 placed.
         Placed:0
         Unplaced:[1 1 1 1 1 1 1 1 1 1]
        >>> PSolution(2)
        Partial QAP solution: [...] with 0 placed.
         Placed:0
         Unplaced:[1 1]
        >>> PSolution()
        Partial QAP solution: None with 0 placed.
         Placed:0
         Unplaced:None
        """
        if n == None:
            pass
        elif sol != None:
            self.n = n
            self.sol = sol
            self.sol_p = <np_int_t*>self.sol.data
            #all facilities start unplaced
            self.unplaced = unplaced
            self.unplaced_p = <np_int_t*>self.unplaced.data
            self.placed = placed
        else:
            self.n = n
            self.sol = np.empty((n), dtype=np.int)
            self.sol_p = <np_int_t*>self.sol.data
            #all facilities start unplaced
            self.unplaced = np.ones((n), dtype=np.int)
            self.unplaced_p = <np_int_t*>self.unplaced.data
            self.placed = 0

    #define properties to set and get various properties of a PSolution

    property sol:
        def __get__(self):
            return self.sol
        def __set__(self, np.ndarray sol):
            self.sol = sol
            self.sol_p = <np_int_t*>sol.data

    property unplaced:
        def __get__(self):
            return self.unplaced
        def __set__(self, np.ndarray unplaced):
            self.unplaced = unplaced
            self.unplaced_p = <np_int_t*>unplaced.data

            
    property placed:
        def __get__(self):
            return self.placed
        def __set__(self, placed):
            self.placed = placed
            
    property n:
        def __get__(self):
            return self.n
        def __set__(self, n):
            self.n = n
   
    def __str__(self):
        '''
        Return a string printout of the object.
        
        Examples:
        
        >>> psol = PSolution(5)
        >>> print psol.__str__()
        Partial QAP solution: [...] with 0 placed.
         Placed:0
         Unplaced:[1 1 1 1 1]
        '''
        s = "Partial QAP solution: "
        s += str(self.sol)
        s += " with %d placed."%self.placed
        s += "\n Placed:%d"%self.placed
        s += "\n Unplaced:"
        s += str(self.unplaced)
        return s
        
    def __repr__(self):
        '''
        Returns a string representation of the object.
        Note: Currently just calls __str__ 
        '''
        return self.__str__()
        
    def __reduce__(self):
        '''
        Reduce the extension into a dictionary for pickling.
        
        Tests(for both __reduce and __setstate__):
        
        >>> from pyQAP import *
        >>> prob = problems.QAPLIB('sko100f')
        >>> psol = heuristics.first(prob)
        >>> psol
        Partial QAP solution: [ 0  1  2 ... 97 98 99] with 100 placed.
         Placed:100
         Unplaced:[0 0 0 ... 0 0 0]
        >>> import pickle, cPickle
        >>> s = pickle.dumps(psol)
        >>> psol_reconstructed = pickle.loads(s)
        >>> psol_reconstructed
        Partial QAP solution: [ 0  1  2 ... 97 98 99] with 100 placed.
         Placed:100
         Unplaced:[0 0 0 ... 0 0 0]
        >>> s = cPickle.dumps(psol)
        >>> psol_reconstructed = cPickle.loads(s)
        >>> psol_reconstructed
        Partial QAP solution: [ 0  1  2 ... 97 98 99] with 100 placed.
         Placed:100
         Unplaced:[0 0 0 ... 0 0 0]
        '''
        return (PSolution, (self.n, None, self.sol, self.unplaced, self.placed))
    
    def change_new(self, index_t facility, PSolution new):
        '''
        This method takes the given facility and creates a new node from the
        given parameter new.
        
        :param facility: the facility to be placed
        :type facility: int
        :param new: the location to put the new partial solution
        :type new: PSolution
        
        Examples:
        
        >>> from pyQAP import *
        >>> psol = PSolution(32)
        >>> new = PSolution(32)
        >>> new
        Partial QAP solution: [...] with 0 placed.
         Placed:0
         Unplaced:[1 1 1 1 1 1 1 1 1 ... 1 1 1 1 1 1 1 1 1 1 1]
        >>> psol.change_new(3,new)
        >>> new
        Partial QAP solution: [ ... 3 ...] with 1 placed.
         Placed:1
         Unplaced:[1 1 1 0 1 1 1 1 1 ... 1 1 1 1 1 1 1 1 1 1 1]
        >>> psol = PSolution(5)
        >>> new = PSolution(5)
        >>> new
        Partial QAP solution: [...] with 0 placed.
         Placed:0
         Unplaced:[1 1 1 1 1]
        >>> psol.change_new(0,new)
        >>> new
        Partial QAP solution: [ ... 0 ...] with 1 placed.
         Placed:1
         Unplaced:[0 1 1 1 1]
        
        '''
        cdef index_t i
        new.placed = self.placed +1
        #copy data from old node into new
        for i in range(self.n):
            new.sol_p[i] = self.sol_p[i]
            new.unplaced_p[i] = self.unplaced_p[i]
        #add the assignment of the new facility
        new.sol_p[self.placed] = facility
        #note that the new facility has been assigned
        new.unplaced_p[facility] = False
    
    def add_new(self, facility):
        """
        Create a new PSolution instance with location placed in the next empty
        location.
        
        :param facility: the facility to be placed
        :type facility: int
        
        Examples:
        
        >>> from pyQAP import *
        >>> psol = PSolution(32)
        >>> psol.add_new(3)
        Partial QAP solution: [ ... 3 ...] with 1 placed.
         Placed:1
         Unplaced:[1 1 1 0 1 1 1 1 1 ... 1 1 1 1 1 1 1 1 1 1 1]
        >>> psol = PSolution(5)
        >>> print psol.add_new(0)
        Partial QAP solution: [... 0 ...] with 1 placed.
         Placed:1
         Unplaced:[0 1 1 1 1]
        """
        #copy with changes into new
        new = PSolution()
        new.n = self.n
        new.placed = self.placed +1
        new.sol = self.sol.copy()
        self.sol_p = <np_int_t*>self.sol.data
        #add the assignment of the new facility
        new.sol[self.placed] = facility
        new.unplaced = self.unplaced.copy()
        self.unplaced_p = <np_int_t*>self.unplaced.data
        #note that the new facility has been assigned
        new.unplaced[facility] = False

        return new
        
    @cython.boundscheck(False)
    @cython.wraparound(False)   
    def cost(self, problem):
        """
        Return the cost of the partial solution using the cost matrices in
        problem.
        
        :param problem: the problem to get the costs from
        :type problem: Problem
        
        Examples:
        
        >>> from pyQAP import *
        >>> from classes import *
        >>> from pyQAP import *
        >>> problem = problems.QAPLIB('had16')
        >>> psol = heuristics.first(problem)
        >>> psol.cost(problem)
        4144.0
        >>> problem = problems.QAPLIB('nug12')
        >>> psol = heuristics.first(problem)
        >>> psol.cost(problem)
        724.0
        >>> problem = problems.QAPLIB('sko100f')
        >>> psol = heuristics.first(problem)
        >>> psol.cost(problem)
        174118.0
        >>> problem = problems.QAPLIB('esc32a')
        >>> psol = heuristics.first(problem)
        >>> psol.cost(problem)
        368.0
        """
        #unpack vars from problem
        cdef np.ndarray[float_t, ndim=2, mode='c'] f = problem.f
        cdef np.ndarray[float_t, ndim=2, mode='c'] d = problem.d
        cdef np.ndarray[float_t, ndim=2, mode='c'] c = problem.c
        cdef np.ndarray[np_int_t, ndim=1, mode='c'] sol = self.sol
        cdef index_t n = self.placed
        cdef index_t i, j
        cdef float_t cost = 0
        for i in range(n):
            for j in range(n):
                cost += f[i,j]*d[sol[i],sol[j]]
        for i in range(n):
            cost += c[i,sol[i]] 
        return cost
        
    def lower_bound(self,problem):
        """
        NOTE: This function caused issues in the parallel solver. It is now
        deprecated. Call vars.lower_bound directly instead.
        
        Returns the lower bound of the current partial solution by calling the 
        bounding method set by set_lower_bound(), stored as current_lower_bound.
        
        :param problem: the problem to get the costs from
        :type problem: Problem
        """
        
        raise DeprecationWarning('This function is deprecated. Call\
                                  vars.lower_bound directly instead')
        
        return vars.lower_bound(self,problem)
    
cdef class Problem:
    """
    A class which contains a quadratic assignment problem. It is constucted
    from the flow (f) and distance (d) matrices and also stores the size of the
    problem as n. C is a linear term used only for sub-problem generation.
    """
    
    cdef np.ndarray f, d, c
    cdef float_t *f_p, *d_p, *c_p
    cdef index_t n
    
    def __init__(self, np.ndarray f, np.ndarray d, np.ndarray c = None):
        '''
        Initialize a new Problem.
        
        Creates a new Problem. If c is not specified, then it is set to the
        appropriate size matrix of all zeros.
        
        :param f: the flow matrix
        :type f: ndarray of doubles
        :param d: the distance matrix
        :type d: ndarray of doubles
        :param c: the constant cost matrix
        :type c: ndarray of doubles
        '''
        self.f = f
        self.d = d
        self.n = f.shape[0]
        if c == None:
            self.c = np.zeros((self.n, self.n), dtype = np.float)
        else:
            self.c = c
        self.f_p = <float_t*>f.data
        self.d_p = <float_t*>d.data
        self.c_p = <float_t*>c.data
        
    property f:
        def __get__(self):
            return self.f
        def __set__(self, np.ndarray f):
            self.f = f
            self.f_p = <float_t*>f.data
            
    property c:
        def __get__(self):
            return self.c
        def __set__(self, np.ndarray c):
            self.c = c
            self.c_p = <float_t*>c.data
            
    property d:
        def __get__(self):
            return self.d
        def __set__(self, np.ndarray d):
            self.d = d
            self.d_p = <float_t*>d.data
            
    property n:
        def __get__(self):
            return self.n
        def __set__(self, n):
            self.f = n
            
    property name:
        def __get__(self):
            return str(self.name)
        def __set__(self, name):
            self.name = name
            
    def __reduce__(self):
        '''
        NOTE: The name of the problem gets mangled with cPickle. This will be
        fixed.
        
        Reduce the extension into a dictionary for pickling.
        
        Tests:
        
        >>> from pyQAP import *
        >>> prob = problems.QAPLIB('sko100f')
        >>> prob
        QAP of size 100 with the flow matrix
        [[  0.   1.   2. ...,  16.  17.  18.]
         [  1.   0.   1. ...,  15.  16.  17.]
         [  2.   1.   0. ...,  14.  15.  16.]
         ..., 
         [ 16.  15.  14. ...,   0.   1.   2.]
         [ 17.  16.  15. ...,   1.   0.   1.]
         [ 18.  17.  16. ...,   2.   1.   0.]]
        and distance matrix
        [[ 0.  5.  1. ...,  4.  0.  0.]
         [ 5.  0.  4. ...,  0.  0.  0.]
         [ 1.  4.  0. ...,  0.  2.  1.]
         ..., 
         [ 4.  0.  0. ...,  0.  3.  1.]
         [ 0.  0.  2. ...,  3.  0.  2.]
         [ 0.  0.  1. ...,  1.  2.  0.]]
        and linear cost matrix
        [[ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]
         ..., 
         [ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]]
        >>> import pickle, cPickle
        >>> s = pickle.dumps(prob)
        >>> prob_reconstructed = pickle.loads(s)
        >>> prob_reconstructed
        QAP of size 100 with the flow matrix
        [[  0.   1.   2. ...,  16.  17.  18.]
         [  1.   0.   1. ...,  15.  16.  17.]
         [  2.   1.   0. ...,  14.  15.  16.]
         ..., 
         [ 16.  15.  14. ...,   0.   1.   2.]
         [ 17.  16.  15. ...,   1.   0.   1.]
         [ 18.  17.  16. ...,   2.   1.   0.]]
        and distance matrix
        [[ 0.  5.  1. ...,  4.  0.  0.]
         [ 5.  0.  4. ...,  0.  0.  0.]
         [ 1.  4.  0. ...,  0.  2.  1.]
         ..., 
         [ 4.  0.  0. ...,  0.  3.  1.]
         [ 0.  0.  2. ...,  3.  0.  2.]
         [ 0.  0.  1. ...,  1.  2.  0.]]
        and linear cost matrix
        [[ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]
         ..., 
         [ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]]
        >>> s = cPickle.dumps(prob)
        >>> prob_reconstructed = cPickle.loads(s)
        >>> prob_reconstructed
        QAP of size 100 with the flow matrix
        [[  0.   1.   2. ...,  16.  17.  18.]
         [  1.   0.   1. ...,  15.  16.  17.]
         [  2.   1.   0. ...,  14.  15.  16.]
         ..., 
         [ 16.  15.  14. ...,   0.   1.   2.]
         [ 17.  16.  15. ...,   1.   0.   1.]
         [ 18.  17.  16. ...,   2.   1.   0.]]
        and distance matrix
        [[ 0.  5.  1. ...,  4.  0.  0.]
         [ 5.  0.  4. ...,  0.  0.  0.]
         [ 1.  4.  0. ...,  0.  2.  1.]
         ..., 
         [ 4.  0.  0. ...,  0.  3.  1.]
         [ 0.  0.  2. ...,  3.  0.  2.]
         [ 0.  0.  1. ...,  1.  2.  0.]]
        and linear cost matrix
        [[ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]
         ..., 
         [ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]
         [ 0.  0.  0. ...,  0.  0.  0.]]
        '''
        return (Problem, (self.f, self.d, self.c))
    
    
    def __str__(self):
        '''
        Returns a string view of the problem.
        
        Examples:
        
        Tested implicitly in other docstrings.
        '''
        s = "QAP of size %d with the flow matrix\n"%(self.n)
        s+= str(self.f)
        s+= "\nand distance matrix\n"
        s+= str(self.d)
        s+= "\nand linear cost matrix\n"
        s+= str(self.c)
        return s
        
    def __repr__(self):
        '''
        Returns a string representation of the problem.
        Note: this currently just calls __str__
        
        Examples:
        
        __str__ is already tested.
        '''
        return self.__str__()
        
    def is_symmetric(self):
        '''
        Check whether the problem is symmetric.
        
        :rtype: boolean
        
        Examples:
        
        >>> from pyQAP import *
        >>> problem = problems.QAPLIB('bur26a')
        >>> problem.is_symmetric()
        False
        >>> problem = problems.QAPLIB('chr12a')
        >>> problem.is_symmetric()
        True
        >>> problem = problems.QAPLIB('els19')
        >>> problem.is_symmetric()
        True
        >>> problem = problems.QAPLIB('esc16a')
        >>> problem.is_symmetric()
        True
        >>> problem = problems.QAPLIB('had20')
        >>> problem.is_symmetric()
        True
        >>> problem = problems.QAPLIB('nug30')
        >>> problem.is_symmetric()
        True
        >>> problem = problems.QAPLIB('lipa90b')
        >>> problem.is_symmetric()
        False
        '''
        #check symmetry of flow matrix
        for i in xrange(self.n):
            for j in xrange(i, self.n):
                if self.f[i, j] != self.f[j, i]:
                    return False
        #check smmetry of distance matrix
        for i in xrange(self.n):
            for j in xrange(i, self.n):
                if self.d[i, j] != self.d[j, i]:
                    return False
        return True
        
class Solution:
    """
    A class for complete solutions which contains the solution array (array),
    the cost of the optimal solution (cost), the number of nodes visited to find
    the optimal solution, and the time taken (time).
    """
    def __init__(self, array, cost, nodes = -1, time = -1):
        '''
        Initialize a new Solution.
        
        This creates a new Solution object. If nodes or time is not specified,
        then they are set to -1. This is done in heuristic solutions where such
        data are meaningless.
        
        :param array: the solution array
        :type array: ndarray of ints
        :param cost: the cost of the Solution
        :type cost: float or double
        :param nodes: the number of nodes explored to find solution
        :type nodes: int
        :param time: the time taked to find the Solution
        :type time: float or double
        '''
        self.array = array
        self.cost = cost
        self.nodes = nodes
        self.time = time
        
    def __str__(self):
        '''
        Returns a string view of the Solution.
        
        Examples:
        
        >>> from pyQAP import *
        >>> Solution([0,1,2,3,4], 10.0, 44, 11)
        QAP Solution: [0, 1, 2, 3, 4] Cost: 10.000000 Nodes visited: 44.000000 Total time: 11.000000
        >>> Solution([4,3,6,7,8,3], 4.36, 678, 2.0)
        QAP Solution: [4, 3, 6, 7, 8, 3] Cost: 4.360000 Nodes visited: 678.000000 Total time: 2.000000
        '''
        s  = "QAP Solution: " + str(self.array)
        s += " Cost: %f Nodes visited: %f Total time: %f"%\
             (self.cost, self.nodes, self.time)
        return s
        
    def __repr__(self):
        '''
        Returns a string representation of the Solution.
        Note: this currently just calls __str__
        
        Examples:
        
        Tested in __str__.
        '''
        return self.__str__()

class ParallelContext:
    '''
    An empty class that is only used to store global information for
    parallel computations.
    '''
    pass
